import colorsys
import os
import copy
import numpy as np
import torch
import torch.nn as nn
from PIL import Image, ImageDraw, ImageFont
from nets.yolo4 import YoloBody
from utils.utils import (DecodeBox, non_max_suppression)
class YOLO(object):
    _defaults = {
        "model_path"  : "logs/Epoch118-Total_Loss2.8951-Val_Loss3.7786.pth",
        "anchors_path": 'model_data/tianchi_anchors.txt',
        "classes_path": 'model_data/tianchi_classes.txt',
        "model_image_size": (416, 416, 3),
        "confidence": 0.5,
        "iou": 0.5,
        "cuda": False
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)
        self.class_names = self._get_class()
        self.anchors = self._get_anchors()
        self.generate()

    def _get_class(self):
        classes_path = os.path.expanduser(self.classes_path)
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    # ---------------------------------------------------#
    #   获得所有的先验框
    # ---------------------------------------------------#
    def _get_anchors(self):
        anchors_path = os.path.expanduser(self.anchors_path)
        with open(anchors_path) as f:
            anchors = f.readline()
        anchors = [float(x) for x in anchors.split(',')]
        return np.array(anchors).reshape([-1, 3, 2])[::-1, :, :]

    # ---------------------------------------------------#
    #   处理斜着照片的
    # ---------------------------------------------------#
    def convertloc(self, x, y, theta, x0, y0):

        # 逆时针
        X = (x - x0) * np.cos(theta) - (y - y0) * np.sin(theta)+x0
        Y = (x - x0) * np.sin(theta) + (y - y0) * np.cos(theta)+y0
        return int(X), int(Y)


    # ---------------------------------------------------#
    #   生成模型
    # ---------------------------------------------------#
    def generate(self):
        # ---------------------------------------------------#
        #   建立yolov4模型
        # ---------------------------------------------------#
        self.net = YoloBody(len(self.anchors[0]), len(self.class_names)).eval()

        # ---------------------------------------------------#
        #   载入yolov4模型的权重
        # ---------------------------------------------------#
        print('Loading weights into state dict...')
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        state_dict = torch.load(self.model_path, map_location=device)
        self.net.load_state_dict(state_dict)
        print('Finished!')

        if self.cuda:
            self.net = nn.DataParallel(self.net)
            self.net = self.net.cuda()

        # ---------------------------------------------------#
        #   建立三个特征层解码用的工具
        # ---------------------------------------------------#
        self.yolo_decodes = []
        for i in range(3):
            self.yolo_decodes.append(
                DecodeBox(self.anchors[i], len(self.class_names), (self.model_image_size[1], self.model_image_size[0])))

        print('{} model, anchors, and classes loaded.'.format(self.model_path))
        # 画框设置不同的颜色
        hsv_tuples = [(x / len(self.class_names), 1., 1.) for x in range(len(self.class_names))]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(
            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors)
        )
    def convert_box(self,box,theta,origin_image ):
        top, left, bottom, right = box
        originx1 = left
        originy1 = top
        originx2 = right
        originy2 = bottom
        theta1 = (theta) * np.pi / 180
        x1, y1 = self.convertloc(originx1, originy1, theta1, origin_image.size[0] / 2, origin_image.size[1] / 2)
        x2, y2 = self.convertloc(originx1, originy2, theta1, origin_image.size[0] / 2, origin_image.size[1] / 2)
        x3, y3 = self.convertloc(originx2, originy1, theta1, origin_image.size[0] / 2, origin_image.size[1] / 2)
        x4, y4 = self.convertloc(originx2, originy2, theta1, origin_image.size[0] / 2, origin_image.size[1] / 2)
        newx1 = np.min([x1, x2, x3, x4])
        newy1 = np.min([y1, y2, y3, y4])
        newx2 = np.max([x1, x2, x3, x4])
        newy2 = np.max([y1, y2, y3, y4])
        box = int(newy1), int(newx1), int(newy2), int(newx2)
        return box


    def detect_image_2(self, image, image_id, result, needrotate=False, num=0):
        theta = 40
        if needrotate:
            theta = 360-20 * num
        origin_image = copy.deepcopy(image)
        image = image.rotate(theta, center=(image.size[0]/2, image.size[1]/2))
        image = image.convert('RGB')
        image_shape = np.array(np.shape(image)[0:2])
        crop_img = image.resize((self.model_image_size[1], self.model_image_size[0]), Image.BICUBIC)
        photo = np.array(crop_img, dtype=np.float32) / 255.0
        photo = np.transpose(photo, (2, 0, 1))
        images = [photo]

        with torch.no_grad():
            images = torch.from_numpy(np.asarray(images))
            if self.cuda:
                images = images.cuda()
            outputs = self.net(images)
            output_list = []
            for i in range(3):
                output_list.append(self.yolo_decodes[i](outputs[i]))
            output = torch.cat(output_list, 1)
            batch_detections = non_max_suppression(output, len(self.class_names), conf_thres=self.confidence,nms_thres=self.iou)
            try:
                batch_detections = batch_detections[0].cpu().numpy()
            except:
                # 在这里处理斜着的照片
                if num >= 19:
                    return origin_image
                num += 1
                needrotate = True
                imagef = self.detect_image_2(origin_image, image_id, result, needrotate, num)
                return imagef

            top_index = batch_detections[:, 4] * batch_detections[:, 5] > self.confidence
            top_label = np.array(batch_detections[top_index, -1], np.int32)
            top_bboxes = np.array(batch_detections[top_index, :4])
            top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(top_bboxes[:, 0], -1), np.expand_dims(
                top_bboxes[:, 1], -1), np.expand_dims(top_bboxes[:, 2], -1), np.expand_dims(top_bboxes[:, 3], -1)
            if len(top_bboxes)<2:
                if num >= 19:
                    return origin_image
                num += 1
                needrotate = True
                imagef = self.detect_image_2(origin_image, image_id, result, needrotate, num)
                return imagef
            top_xmin = top_xmin / self.model_image_size[1] * image_shape[1]
            top_ymin = top_ymin / self.model_image_size[0] * image_shape[0]
            top_xmax = top_xmax / self.model_image_size[1] * image_shape[1]
            top_ymax = top_ymax / self.model_image_size[0] * image_shape[0]
            boxes = np.concatenate([top_ymin, top_xmin, top_ymax, top_xmax], axis=-1)
        for i in range(len(boxes)):
            top, left, bottom, right = boxes[i]
            originx1 = left
            originy1 = top
            originx2 = right
            originy2 = bottom
            theta1 = (theta)*np.pi/180
            x1, y1 = self.convertloc(originx1, originy1, theta1, origin_image.size[0]/2, origin_image.size[1]/2)
            x2, y2 = self.convertloc(originx1, originy2, theta1, origin_image.size[0]/2, origin_image.size[1]/2)
            x3, y3 = self.convertloc(originx2, originy1, theta1, origin_image.size[0]/2, origin_image.size[1]/2)
            x4, y4 = self.convertloc(originx2, originy2, theta1, origin_image.size[0]/2, origin_image.size[1]/2)
            newx1 = np.min([x1, x2, x3, x4])
            newy1 = np.min([y1, y2, y3, y4])
            newx2 = np.max([x1, x2, x3, x4])
            newy2 = np.max([y1, y2, y3, y4])
            boxes[i] = newy1, newx1, newy2, newx2
        peopleArr = []
        badgeArr = []
        clothesArr = []
        wrongclothesArr = []
        for i, c in enumerate(top_label):
            """
                person
                badge
                wrongclothes
                clothes
            """
            item = {}
            item["image_id"] = image_id
            predicted_class = self.class_names[c]
            item["score"] = 1
            top, left, bottom, right = boxes[i]
            top = max(0, np.floor(top + 0.5).astype('int32'))
            left = max(0, np.floor(left + 0.5).astype('int32'))
            bottom = min(np.shape(image)[0], np.floor(bottom + 0.5).astype('int32'))
            right = min(np.shape(image)[1], np.floor(right + 0.5).astype('int32'))
            item["bbox"] = [int(left), int(top), int(right), int(bottom)]

            newtop = max(0, np.floor(top + 0.5).astype('int32'))
            newleft = max(0, np.floor(left + 0.5).astype('int32'))
            newbottom = min(np.shape(image)[0], np.floor(bottom + 0.5).astype('int32'))
            newright = min(np.shape(image)[1], np.floor(right + 0.5).astype('int32'))
            draw = ImageDraw.Draw(origin_image)
            for i in range(2):
                draw.rectangle([newleft+i, newtop+i, newright+i, newbottom+i], outline=self.colors[self.class_names.index(predicted_class)])
            font = ImageFont.truetype(font='model_data/simhei.ttf',
                                      size=np.floor(3e-2 * np.shape(origin_image)[1] + 0.5).astype('int32'))
            draw.text([newleft, newtop], str(predicted_class), fill=(248, 59, 22), font=font)
            del draw
            if c == 0:
                peopleArr.append(item)
            elif c == 1:
                badgeArr.append(item)
            elif c == 2:
                wrongclothesArr.append(item)
            else:
                clothesArr.append(item)
        print(f"第{num + 1}次Image_id:{image_id}  labels:{top_label}")
        for item1 in badgeArr:
            item1['used'] = False
        for item1 in clothesArr:
            item1['used'] = False

        for item0 in peopleArr:
            flag = 0
            left0, top0, right0, bottom0 = item0["bbox"]
            for item1 in badgeArr:
                left1, top1, right1, bottom1 = item1["bbox"]
                length = (right1 - left1) / 5
                if left0-length <= left1 and top0-length<= top1 and right0+length>= right1 and (top1+bottom1)/2<(top0/4+3*bottom0/4):
                    flag = 1
                    item1['used']=True
                    itemResult = copy.deepcopy(item0)
                    itemResult["category_id"] = 1
                    # box = itemResult['bbox']
                    # itemResult['bbox'] = box#self.convert_box(box, theta, origin_image)
                    # itemResult['bbox'] = self.decrease_box(itemResult['bbox'])
                    result.append(itemResult)
                else:
                    pass
            for item1 in clothesArr:
                left1, top1, right1, bottom1 = item1["bbox"]
                length = (right1-left1)/4
                if left0 - length <= left1 and top0 - length <= top1 and right0 + length >= right1:
                    item1['used'] = True
                    flag = 1
                    itemResult = copy.deepcopy(item0)
                    itemResult["category_id"] = 2
                    # box = itemResult['bbox']
                    # itemResult['bbox'] = box#self.convert_box(box, theta, origin_image)
                    # itemResult['bbox'] = self.decrease_box(itemResult['bbox'])
                    result.append(itemResult)
                else:
                    pass

            if flag == 0:
                itemResult = copy.deepcopy(item0)
                # 改变坐标
                # box = itemResult['bbox']
                # itemResult['bbox'] = self.convert_box(box, theta, origin_image)
                # itemResult['bbox'] = self.decrease_box(itemResult['bbox'])
                itemResult["category_id"] = 3
                result.append(itemResult)
        for item in clothesArr:
            if not item['used']:
                itemResult = {}
                itemResult["category_id"] = 2
                itemResult["image_id"] = image_id
                itemResult["score"] = 1
                left, top, right, bottom = item['bbox']
                length = (bottom - top) / 5
                bottom = bottom
                top = max(0, top - length)
                itemResult['bbox'] = [int(left), int(top), int(right), int(bottom)]
                # box = itemResult['bbox']
                # itemResult['bbox'] = self.convert_box(box, theta, origin_image)
                # itemResult['bbox'] = self.decrease_box(itemResult['bbox'])
                result.append(itemResult)
        for item in badgeArr:
            if not item['used']:
                itemResult = {}
                itemResult["category_id"] = 1
                itemResult["image_id"] = image_id
                itemResult["score"] = 1
                left, top, right, bottom = item['bbox']
                length = (bottom - top)
                leng = right - left
                left = max(0, left - leng)
                right = min(origin_image.size[0], right + leng*2)
                bottom = min(origin_image.size[1], bottom + length*6)
                top = max(0, top - length*3)
                itemResult['bbox'] = [int(left), int(top), int(right), int(bottom)]
                # box = itemResult['bbox']
                # itemResult['bbox'] = self.convert_box(box, theta, origin_image)
                # itemResult['bbox'] = self.decrease_box(itemResult['bbox'])
                result.append(itemResult)

        return origin_image


    # ---------------------------------------------------#
    #   检测图片
    # ---------------------------------------------------#
    def detect_image(self, image, image_id, result):
        origin_image = copy.deepcopy(image)
        image = image.convert('RGB')
        image_shape = np.array(np.shape(image)[0:2])
        crop_img = image.resize((self.model_image_size[1], self.model_image_size[0]), Image.BICUBIC)
        photo = np.array(crop_img, dtype=np.float32) / 255.0
        photo = np.transpose(photo, (2, 0, 1))
        images = [photo]
        with torch.no_grad():
            images = torch.from_numpy(np.asarray(images))
            if self.cuda:
                images = images.cuda()
            outputs = self.net(images)
            output_list = []
            for i in range(3):
                output_list.append(self.yolo_decodes[i](outputs[i]))
            output = torch.cat(output_list, 1)
            batch_detections = non_max_suppression(output, len(self.class_names), conf_thres=self.confidence, nms_thres=self.iou)
            try:
                batch_detections = batch_detections[0].cpu().numpy()
            except:
                # 在这里处理斜着的照片
                image = self.detect_image_2(origin_image, image_id, result, False, 0)
                return image

            top_index = batch_detections[:, 4] * batch_detections[:, 5] > self.confidence
            top_label = np.array(batch_detections[top_index, -1], np.int32)
            top_bboxes = np.array(batch_detections[top_index, :4])
            if len(top_bboxes)<2:
                imagef = self.detect_image_2(origin_image, image_id, result, False, 0)
                return imagef
            top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(top_bboxes[:, 0], -1), np.expand_dims(
                top_bboxes[:, 1], -1), np.expand_dims(top_bboxes[:, 2], -1), np.expand_dims(top_bboxes[:, 3], -1)

            top_xmin = top_xmin / self.model_image_size[1] * image_shape[1]
            top_ymin = top_ymin / self.model_image_size[0] * image_shape[0]
            top_xmax = top_xmax / self.model_image_size[1] * image_shape[1]
            top_ymax = top_ymax / self.model_image_size[0] * image_shape[0]
            boxes = np.concatenate([top_ymin, top_xmin, top_ymax, top_xmax], axis=-1)

        peopleArr = []
        badgeArr = []
        clothesArr = []
        wrongclothesArr = []
        print(f"Image_id:{image_id}  labels:{top_label}")

        for i, c in enumerate(top_label):
            """
            person
            badge
            wrongclothes
            clothes
            """
            item = {}
            item["image_id"] = image_id
            predicted_class = self.class_names[c]
            item["score"] = 1
            top, left, bottom, right = boxes[i]
            top = max(0, np.floor(top + 0.5).astype('int32'))
            left = max(0, np.floor(left + 0.5).astype('int32'))
            bottom = min(np.shape(image)[0], np.floor(bottom + 0.5).astype('int32'))
            right = min(np.shape(image)[1], np.floor(right + 0.5).astype('int32'))
            item["bbox"] = [int(left), int(top), int(right), int(bottom)]
            draw = ImageDraw.Draw(image)
            for i in range(1):
                draw.rectangle([left+i, top+i, right+i, bottom+i], outline=self.colors[self.class_names.index(predicted_class)])
            font = ImageFont.truetype(font='model_data/simhei.ttf',
                                      size=np.floor(3e-2 * np.shape(image)[1] + 0.5).astype('int32'))
            draw.text([left, top], str(predicted_class), fill=(248, 59, 22), font=font)
            del draw
            if c == 0:
                peopleArr.append(item)
            elif c == 1:
                badgeArr.append(item)
            elif c == 2:
                wrongclothesArr.append(item)
            else:
                clothesArr.append(item)
        for item1 in badgeArr:
            item1['used'] = False
        for item1 in clothesArr:
            item1['used'] = False
        for item0 in peopleArr:
            flag = 0
            left0, top0, right0, bottom0 = item0["bbox"]
            for item1 in badgeArr:
                left1, top1, right1, bottom1 = item1["bbox"]
                length = (right1 - left1) / 5#一定要大于3,小于等于5
                if left0-length <= left1 and top0-length<= top1 and right0+length>= right1 and (top1+bottom1)/2<(top0/4+3*bottom0/4):
                    flag = 1
                    item1['used']=True
                    itemResult = copy.deepcopy(item0)
                    itemResult["category_id"] = 1
                    result.append(itemResult)
                else:
                    pass
            for item1 in clothesArr:
                left1, top1, right1, bottom1 = item1["bbox"]
                length = (right1-left1)/4
                if max(left0 - length, 0) <= left1 and max(0, top0 - length) <= top1 and min(origin_image.size[0], right0 + length) >= right1:
                    item1['used'] = True
                    flag = 1
                    itemResult = copy.deepcopy(item0)
                    itemResult["category_id"] = 2
                    result.append(itemResult)
                else:
                    pass
            if flag == 0:
                itemResult = copy.deepcopy(item0)
                itemResult["category_id"] = 3
                result.append(itemResult)
        for item in clothesArr:
            if not item['used']:
                itemResult = {}
                itemResult["category_id"] = 2
                itemResult["image_id"] = image_id
                itemResult["score"] = 1
                left, top, right, bottom = item['bbox']
                length = (bottom-top)/5
                bottom = bottom
                top = max(0, top-length)
                itemResult['bbox'] = [int(left), int(top), int(right), int(bottom)]
                result.append(itemResult)

        for item in badgeArr:
            if not item['used']:
                itemResult={}
                itemResult["category_id"] = 1
                itemResult["image_id"] = image_id
                itemResult["score"] = 1
                left, top, right, bottom = item['bbox']
                length = (bottom-top)
                leng = (right-left)
                left = max(left-leng, 0)
                right = min(right+leng*2, origin_image.size[0])
                bottom = min(bottom+length*6, origin_image.size[1])
                top = max(top-length*3, 0)
                itemResult['bbox'] = [int(left), int(top), int(right), int(bottom)]
                result.append(itemResult)
        return image

